#! /usr/bin/python
# encoding=utf-8
# Created by Fenglu Niu on 2025/5/6 11:39
import os

import requests
from dotenv import load_dotenv
from fengluU import exec_time


class Dify(object):
    __instance = None
    __is_first = True

    def __new__(cls):
        if not cls.__instance:
            cls.__instance = object.__new__(cls)
        return cls.__instance

    def __init__(self):
        if Dify.__is_first:
            Dify.__is_first = False
            load_dotenv()
            self.__api_base = os.getenv('DIFY_API_BASE')
            self.__dataset_api_key = os.getenv('DIFY_DATASET_API_KEY')

    @exec_time
    def retrieve(self, query: str, dataset_id: str):
        dataset_url = f'{self.__api_base}/datasets/{dataset_id}/retrieve'
        body = {'query': query,
                'retrieval_model': {'search_method': 'hybrid_search', 'reranking_enable': 'false', 'top_k': 2,
                                    'score_threshold_enabled': True, 'score_threshold': 0.55}}
        resp = requests.post(dataset_url, json=body,
                             headers={'Authorization': f'Bearer {self.__dataset_api_key}'}).json()
        records = resp['records']
        result = ''
        for record in records:
            result += record['segment']['content'] + '\n'
        return result


if __name__ == '__main__':
    dify = Dify()
    print(dify.retrieve('明天有什么菜啊', 'df4fcd65-302c-47d6-8520-817256b55266'))
